import struct
import random
import hashlib
import uuid
import socket
from unittest import TestCase

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes, hmac
from cryptography.hazmat.primitives.asymmetric import ec, padding
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
from cryptography import x509

from ssl.common import generate_gmt_time, derive_keys
from sslserver import read_ecc_private_key_pem




# 生成 ClientHello 报文
def create_client_hello(hostname="example.com"):
    handshake_type = b'\x01'  # ClientHello
    client_version = b'\x03\x03'
    random_bytes = generate_gmt_time()

    session_id_length = b'\x00'
    session_id = b''

    cipher_suites = [
        0xC0, 0x30,  # ECDHE-RSA-AES256-GCM-SHA384
        0xC0, 0x2F,  # ECDHE-RSA-AES128-GCM-SHA256
        0x00, 0x9C  # AES256-GCM-SHA384
    ]
    cipher_suites_length = struct.pack('>H', len(cipher_suites))

    compression_methods = [0x00]  # NULL compression
    compression_methods_length = struct.pack('>B', len(compression_methods))

    extensions = []

    sni_hostname = hostname.encode('ascii')
    sni_server_name = [
        struct.pack('>H', 3 + len(sni_hostname)),  # Server name list length
        b'\x00',  # Server name type
        struct.pack('>H', len(sni_hostname)),  # Hostname length
        sni_hostname  # Hostname bytes
    ]
    sni_server_name = b''.join(sni_server_name)
    sni_extension = [
        b'\x00\x00',  # Type: server_name (0)
        struct.pack('>H', len(sni_server_name)),  # Length
        sni_server_name
    ]
    extensions.extend(sni_extension)

    # elliptic_curves = [
    #     b'\x00\x0A',  # Type: elliptic_curves (10)
    #     struct.pack('>H', 4),  # Length
    #     struct.pack('>H', 0x0017),  # secp256r1
    #     struct.pack('>H', 0x0018)  # secp384r1
    # ]
    # extensions.extend(elliptic_curves)
    extensions = b''.join(extensions)
    extensions_length = struct.pack('>H', len(extensions))

    client_hello = (
            client_version +
            random_bytes +
            session_id_length +
            session_id +
            cipher_suites_length +
            bytes(cipher_suites) +
            compression_methods_length +
            bytes(compression_methods) +
            extensions_length +
            extensions
    )
    handshake_length = struct.pack('>I', len(client_hello))[1:]  # 取低3位字节，错了就得花时间排查
    client_hello = (handshake_type + handshake_length + client_hello)
    return wrap_tls(client_hello)


# 生成 ClientKeyExchange 报文
def create_client_key_exchange():
    server_private_key = read_ecc_private_key_pem("./ecc-key.pem")
    server_public_key = server_private_key.public_key()
    private_key = read_ecc_private_key_pem('./client-ecc-key.pem')
    public_key = private_key.public_key()
    shared_secret = private_key.exchange(ec.ECDH(), server_public_key)
    master_secret = HKDF(
        algorithm=hashes.SHA256(),
        length=48,
        salt=None,
        info=b'tls12 master secret',
    ).derive(shared_secret)

    der_public = public_key.public_bytes(encoding=Encoding.X962,format=PublicFormat.CompressedPoint)

    pre_master_secret = bytes(random.getrandbits(8) for _ in range(48))



    handshake = (
            b'\x10' +  # Handshake type: ClientKeyExchange
            struct.pack('>I', len(der_public) + 1)[1:] +
            struct.pack('>H', len(der_public))[1:] +
            der_public
    )
    return wrap_tls(handshake), master_secret


# 生成 ChangeCipherSpec 报文
def create_change_cipher_spec():
    record = b'\x14\x03\x03\x00\x01\x01'
    return record


# 生成 Finished 报文
def create_finished(handshake_messages, master_secret, is_client=True):
    label = b'client finished' if is_client else b'server finished'
    verify_data = HKDF(
        algorithm=hashes.SHA256(),
        length=12,
        salt=b'',
        info=label,
    ).derive(master_secret)

    handshake = (
            b'\x14' +  # Handshake type: Finished
            struct.pack('>I', 12)[1:] +
            verify_data
    )
    return wrap_tls(handshake)


# 封装握手消息到 TLS 记录层
def wrap_tls(handshake_data):
    record_type = b'\x16'
    protocol_version = b'\x03\x03'
    record_length = struct.pack('>H', len(handshake_data))
    return record_type + protocol_version + record_length + handshake_data

# 加密数据
def encrypt_data(data, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    encryptor = cipher.encryptor()
    padded_data = data + b'\x00' * ((16 - len(data) % 16) % 16)
    ciphertext = encryptor.update(padded_data) + encryptor.finalize()
    return ciphertext


# 解密数据
def decrypt_data(ciphertext, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    decryptor = cipher.decryptor()
    padded_data = decryptor.update(ciphertext) + decryptor.finalize()
    return padded_data.rstrip(b'\x00')


# 发送加密数据
def send_encrypted_data(sock, data, key, iv):
    encrypted = encrypt_data(data, key, iv)
    record = (
            b'\x17' +  # Application data
            b'\x03\x03' +  # TLS 1.2
            struct.pack('>H', len(encrypted)) +
            encrypted
    )
    sock.send(record)


# 接收加密数据
def receive_encrypted_data(sock, key, iv):
    record_type = sock.recv(1)
    if record_type != b'\x17':
        raise ValueError("Not application data")

    version = sock.recv(2)
    length = struct.unpack('>H', sock.recv(2))[0]
    ciphertext = sock.recv(length)
    return decrypt_data(ciphertext, key, iv)


def read_file_to_string(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
            print(f"成功读取文件 {file_path} 的内容")
            return content
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到")
    except Exception as e:
        print(f"读取文件时出现错误: {e}")
    return None


def client():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.connect(("127.0.0.1", 8888))

        client_hello = create_client_hello()
        sock.send(client_hello)

        server_hello = sock.recv(4096)
        certificate = sock.recv(4096)
        server_key_exchange = sock.recv(4096)
        server_hello_done = sock.recv(4096)

        cert = x509.load_der_x509_certificate(certificate[15:])
        server_public_key = cert.public_key()

        client_key_exchange, master_secret = create_client_key_exchange()
        sock.send(client_key_exchange)

        change_cipher_spec = create_change_cipher_spec()
        sock.send(change_cipher_spec)

        handshake_messages = client_hello + client_key_exchange
        client_finished = create_finished(handshake_messages, master_secret, is_client=True)
        sock.send(client_finished)

        server_change_cipher = sock.recv(4096)
        server_finished = sock.recv(4096)

        client_random = client_hello[6:38]
        server_random = server_hello[6:38]
        client_write_key, server_write_key, client_write_iv, server_write_iv = derive_keys(
            master_secret, client_random, server_random
        )

        while True:
            message = input("Enter a message to send (or 'quit' to exit): ")
            if message.lower() == 'quit':
                break
            send_encrypted_data(sock, message.encode(), client_write_key, client_write_iv)
            try:
                response = receive_encrypted_data(sock, server_write_key, server_write_iv)
                print(f"Received from server: {response.decode()}")
            except ValueError as e:
                print(f"Error: {e}")


def parse_extensions(data):
    extensions = []
    offset = 0
    while offset < len(data):
        # 解析 extension_type
        extension_type = struct.unpack('!H', data[offset:offset + 2])[0]
        offset += 2
        # 解析 extension_data 的长度
        data_length = struct.unpack('!H', data[offset:offset + 2])[0]
        offset += 2
        # 解析 extension_data
        extension_data = data[offset:offset + data_length]
        offset += data_length
        extensions.append((extension_type, extension_data))
    return extensions


class Test(TestCase):
    def test_create_client_hello(self):
        create_client_hello()

    def test_parse_extensions(self):
        # 示例数据
        clienthello_extensions_data = b'\x00\x0A\x00\x05\x00\x03\x00\x00\x00'
        extensions = parse_extensions(clienthello_extensions_data)
        for ext_type, ext_data in extensions:
            print(f"Extension Type: 0x{ext_type:04x}, Extension Data: {ext_data.hex()}")


if __name__ == "__main__":
    client()
